package zuo.algo.c18;

import zuo.algo.data.Node;
import zuo.algo.util.SortUtils;
import zuo.algo.util.TreeUtils;

public class AVL<T extends Comparable<T>> {

    private Node<T> root;

    public void add(T value) {
        if (value == null) {
            return;
        }
        root = add(root, value);
    }

    private Node<T> add(Node<T> node, T value) {
        if (node == null) {
            return new Node<>(value, 1);
        }
        int cmp = value.compareTo(node.value);
        if (cmp < 0) {
            node.left = add(node.left, value);
        } else if (cmp > 0) {
            node.right = add(node.right, value);
        }
        node = balance(node);
        refreshHeight(node);
        return node;
    }

    private Node<T> balance(Node<T> node) {
        if (Math.abs(height(node.left) - height(node.right)) <= 1) {
            return node;
        }
        // LL型
        //          5                       3
        //         / \                     / \
        //        3   6        右旋        2   5
        //       / \          ---->      /   / \
        //      2   4                   1   4   6
        //     /
        //    1
        // L型：删除时可能出现
        //          6                          6                 3
        //         / \                        / \               / \
        //        3   7       删除7或8        3   7     右旋     2   6
        //       / \   \    ----------->    / \       ---->   /   / \
        //      2   4   8                  2   4             1   4   7
        //     /     \                    /     \                 \
        //    1       5                  1       5                 5
        if (height(node.left) > height(node.right) &&
                height(node.left.left) >= height(node.left.right)) {
            return rightHead(node);
        }
        // LR型
        //          5                          5                       3
        //         / \                        / \                     / \
        //        2   6       2左旋           3   6        右旋        2   5
        //       / \         ------>        / \          ---->      /   / \
        //      1   3                      2   4                   1   4   6
        //           \                    /
        //            4                  1
        if (height(node.left) > height(node.right) &&
                height(node.left.right) > height(node.left.left)) {
            node.left = leftHead(node.left);
            return rightHead(node);
        }
        // RR型
        //          5                        7
        //         / \                      / \
        //        3   7         左旋        5   8
        //           / \       ---->      / \   \
        //          6   8                3   6   9
        //               \
        //                9
        // R型：删除时可能出现
        //          4                          4                 7
        //         / \                        / \               / \
        //        3   7       删除2或3        3   7     左旋    4   8
        //       /   / \    ----------->        / \   ----->  / \   \
        //      2   6   8                      6   8         3   6   9
        //         /     \                    /     \           /
        //        5       9                  5       9         5
        if (height(node.right) > height(node.left) &&
                height(node.right.right) >= height(node.right.left)) {
            return leftHead(node);
        }
        // RR型
        //          5                         5                        7
        //         / \                       / \                      / \
        //        3   8         8右旋        3   7         左旋        5   8
        //           / \       ------>         / \       ----->      / \   \
        //          7   9                     6   8                 3   6   9
        //         /                               \
        //        6                                 9
        if (height(node.right) > height(node.left) &&
                height(node.right.left) > height(node.right.right)) {
            node.right = rightHead(node.right);
            return leftHead(node);
        }
        return node;
    }

    /**
     * 左旋 - 头倒向左边
     */
    private Node<T> leftHead(Node<T> node) {
        Node<T> node1 = node.right;
        node.right = node1.left;
        node1.left = node;
        refreshHeight(node);
        return node1;
    }

    /**
     * 右旋 - 头倒向右边
     */
    private Node<T> rightHead(Node<T> node) {
        Node<T> node1 = node.left;
        node.left = node1.right;
        node1.right = node;
        refreshHeight(node);
        return node1;
    }

    private int height(Node<T> node) {
        if (node == null) {
            return 0;
        }
        return node.n;
    }

    private void refreshHeight(Node<T> node) {
        if (node == null) {
            return;
        }
        node.n = Math.max(height(node.left), height(node.right)) + 1;
    }

    public void delete(T value) {
        if (value == null) {
            return;
        }
        root = delete(root, value);
    }

    private Node<T> delete(Node<T> node, T value) {
        if (node == null) {
            return null;
        }
        int cmp = value.compareTo(node.value);
        if (cmp < 0) {
            node.left = delete(node.left, value);
        } else if (cmp > 0) {
            node.right = delete(node.right, value);
        } else {
            if (node.left == null) {
                return node.right;
            }
            if (node.right == null) {
                return node.left;
            }
            // 获取右子树的最小值替换当前节点
            Node<T> delNode = node;
            node = min(delNode.right);
            node.right = resizeRight(delNode.right);
            node.left = delNode.left;
        }
        node = balance(node);
        refreshHeight(node);
        return node;
    }

    private Node<T> resizeRight(Node<T> node) {
        if (node.left == null) {
            return node.right;
        }
        node.left = resizeRight(node.left);
        node = balance(node);
        refreshHeight(node);
        return node;
    }

    private Node<T> min(Node<T> node) {
        if (node.left == null) {
            return node;
        }
        return min(node.left);
    }

    public static void main(String[] args) {
        AVL<Integer> avl = new AVL<>();
//        avl.add(40);
//        avl.add(20);
//        avl.add(30);
//        avl.add(10);
//        avl.add(60);
//        avl.add(50);
//        avl.add(70);
//
//        TreeUtils.printTree(avl.root);
//
//        avl.add(25);
//        avl.add(27);
//        avl.add(22);
//        avl.add(23);
//
//        TreeUtils.printTree(avl.root);
//
//        avl.delete(20);
//        TreeUtils.printTree(avl.root);
//        int[] arr = new int[] {1640, 8670, 9655, 9584, 78, 5879, 1885};
        for (int i = 0; i < 100000000; i++) {
            int[] arr = SortUtils.getRandomArr(1000, 1000000);
            for (int num : arr) {
                avl.add(num);
                boolean balance = TreeUtils.isBalance(avl.root);
                if (!balance) {
                    throw new RuntimeException();
                }
            }
            for (int num : arr) {
                avl.delete(num);
                boolean balance = TreeUtils.isBalance(avl.root);
                if (!balance) {
                    throw new RuntimeException();
                }
            }
            System.out.println(i);
        }
        System.out.println();
    }
}
